import numpy as np
import os
from tqdm import trange


class ssdiff(object):

    def __init__(self, **kwargs):

        self.wd = kwargs.get('wd', os.getcwd())
        self.filename = kwargs.get('filename', None)
        self.fmt = kwargs.get('fmt', 'csv')
        self.ssArray = kwargs.get('ss', None)
        self.numIter = self.ssArray.shape[0]

    def calcKASD(self, **kwargs):
        """
        计算 Kernel Average Slip Deviation
        :param kwargs:
        wltype: worklist type; 1: f=open(*) or f.readlines; 2: numpy array
        :return:
        """
        workList = kwargs.get('worklist', None)
        wlt = kwargs.get('wltype', 1)

        if wlt == 1:
            if not isinstance(workList, list):
                workList = workList.readlines()
        else:
            pass

        kasdList = np.array(())
        for i in trange(self.numIter):
            if wlt == 1:
                tempWorkList = np.fromstring(workList[i], dtype=int, sep='\t')
            else:
                tempWorkList = workList[i]

            a = self.ssArray[i]
            b = self.ssArray[tempWorkList - 1]
            diff = np.abs(b - a)
            diff = np.mean(diff)
            kasdList = np.append(kasdList, diff)

        np.savetxt(os.path.join(self.wd, '{}_KASD.{}'.format(self.filename, self.fmt)), kasdList,
                   fmt='%3f', delimiter='\t')
        np.savetxt(os.path.join(self.wd, '{}_KASD_MEAN.{}'.format(self.filename, self.fmt)),
                   kasdList / np.mean(kasdList, axis=0), fmt='%3f', delimiter='\t')

    def calcGASD(self, **kwargs):
        """
        计算 Grain Average Slip Deviation
        :param kwargs:
        :return:
        """
        workList = kwargs.get('worklist', None)
        if not isinstance(workList, list):
            workList = workList.readlines()

        gasdList = np.array(())
        for i in trange(self.numIter):
            tempWorkList = np.fromstring(workList[i], dtype=int, sep='\t')
            a = self.ssArray[i]
            b = self.ssArray[tempWorkList - 1]
            diff = np.abs(b - a)
            diff = np.mean(diff)
            gasdList = np.append(gasdList, diff)

        np.savetxt(os.path.join(self.wd, '{}_GASD.{}'.format(self.filename, self.fmt)), gasdList,
                   fmt='%3f', delimiter='\t')
        np.savetxt(os.path.join(self.wd, '{}_GASD_MEAN.{}'.format(self.filename, self.fmt)),
                   gasdList / np.mean(gasdList), fmt='%3f', delimiter='\t')

    def calcGSS(self, **kwargs):
        """
        计算 Grain Slip Spread
        :param kwargs:
        wl: 是否有worklist；如果为False， 将通过 ELESET 数组来计算
        :return:
        """
        elesetArray = kwargs.get('eleset', None)
        wl = kwargs.get('wl', False)

        if wl:
            workList = kwargs.get('worklist', None)
            if not isinstance(workList, list):
                workList = workList.readlines()

        elesetNum = np.unique(elesetArray)
        numEleset = elesetNum.shape[0]

        gssList = np.zeros(self.numIter)
        for i in trange(numEleset):

            if not wl:
                tempWorkList = np.where(elesetArray == elesetNum[i])[0]
            else:
                tempWorkList = np.fromstring(workList[i], dtype=int, sep='\t')

            tempSSArray = self.ssArray[tempWorkList]
            meanSS = np.mean(np.abs(tempSSArray))
            diff = np.mean(np.abs(tempSSArray - meanSS))
            gssList[workList] = diff

        np.savetxt(os.path.join(self.wd, '{}_GSS.{}'.format(self.filename, self.fmt)), gssList,
                   fmt='%3f', delimiter='\t')
        np.savetxt(os.path.join(self.wd, '{}_GSS_MEAN.{}'.format(self.filename, self.fmt)),
                   gssList / np.mean(gssList), fmt='%3f', delimiter='\t')
